import type { ChatInit, ChatStatus } from "ai"
import { AbstractChat } from "ai"

import type { BizUIMessage } from "../types"
import { ZustandChatState } from "./chat-state"
import type { ChatSlice } from "./types"

// Custom Chat class that uses Zustand-integrated state
export class ZustandChat extends AbstractChat<BizUIMessage> {
  override state: ZustandChatState
  #unsubscribeFns: (() => void)[] = []

  constructor(
    { messages, ...init }: ChatInit<BizUIMessage>,
    updateZustandState: (updater: (state: ChatSlice) => ChatSlice) => void,
  ) {
    const state = new ZustandChatState(messages, updateZustandState, init.id || "")
    super({ ...init, state })
    this.state = state

    const baseResumeStream = this.resumeStream.bind(this)
    // Track resume calls so the state can ignore the temporary "submitted" status when no stream exists.
    this.resumeStream = async (...args) => {
      this.state.setResumingStream(true)
      try {
        return await baseResumeStream(...args)
      } finally {
        this.state.setResumingStream(false)
      }
    }
  }

  // Public getter for state access
  get chatState() {
    return this.state
  }

  // Cleanup method
  async destroy(): Promise<void> {
    await this.stop()
    // Unsubscribe from AI SDK callbacks
    this.#unsubscribeFns.forEach((unsubscribe) => unsubscribe())
    this.#unsubscribeFns = []

    this.state.destroy()
  }

  protected override setStatus({ status, error }: { status: ChatStatus; error?: Error }): void {
    super.setStatus({ status, error })
    this.state.status = status
    this.state.error = error
  }
}
